# encoding=utf-8
import argparse
import hashlib
import json
import os
import tempfile
import time
import logging
import traceback
import fastapi
import uvicorn
from asyncio import Lock
from fastapi import FastAPI
from starlette.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from logging.handlers import TimedRotatingFileHandler
from DenseGradeWrapper import DensenetGrade

parser = argparse.ArgumentParser()
parser.add_argument("--path", type=str)
args = parser.parse_args()
if args.path is None:
    logging.error("run path is empty")
    exit(-1)

run_path = args.path

os.makedirs("runtime-cache", exist_ok=True)
os.makedirs("runtime-log", exist_ok=True)


def set_logger():
    logger = logging.getLogger()  # root
    logger.setLevel(logging.INFO)

    file_handler = TimedRotatingFileHandler(
        filename=os.path.join("runtime-log", "app-log"),
        when='D', interval=1
    )
    file_handler.suffix = "%Y-%m-%d.log"

    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)

    logging_format = logging.Formatter(
        "%(asctime)s - %(filename)s - %(funcName)s - %(lineno)d - %(levelname)s - %(message)s")

    file_handler.setFormatter(logging_format)
    console_handler.setFormatter(logging_format)

    logger.addHandler(file_handler)
    logger.addHandler(console_handler)


set_logger()
dense_grade_model = DensenetGrade(run_path)


def generate_time_mills():
    return int(round(time.time() * 1000))


def generate_task_id(name: str):
    str_data = name + str(generate_time_mills())
    encoded = hashlib.sha256(str_data.encode()).hexdigest()
    return encoded


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"]
)

async_lock = Lock()


class DenseGradeController:
    @staticmethod
    @app.post("/run/dense121/five")
    async def run(file: fastapi.UploadFile):
        # save to tmp, use PIL to decode
        temp_dir_name = generate_task_id("hello")
        abs_workdir = os.path.join("runtime-cache", temp_dir_name)
        abs_user_path = os.path.join(abs_workdir, "user_input")
        os.makedirs(abs_user_path, exist_ok=True)

        logging.info("save user path: {}".format(abs_user_path))

        temp_file = os.path.basename(tempfile.NamedTemporaryFile(suffix=".png").name)
        abs_file_path = os.path.join(abs_user_path, temp_file)
        with open(abs_file_path, 'wb') as abs_file:
            abs_file.write(await file.read())

        logging.info("write done file: {}".format(abs_file_path))

        # run infer
        async with async_lock:
            try:
                start_time = time.perf_counter()
                # get possible
                out = dense_grade_model.infer(abs_file_path)

                # save to local for debug
                out_path = os.path.join(abs_workdir, "out_raw_data")
                os.makedirs(out_path, exist_ok=True)

                # txt file
                txt_temp_file = temp_file.replace(".png", ".txt")
                with open(os.path.join(out_path, txt_temp_file), "w") as abs_txt_file:
                    abs_txt_file.write(json.dumps(out))

                # print time cost
                end_time = time.perf_counter()
                logging.info("infer cost: {}s".format(end_time - start_time))

                # return
                return out
            except Exception as e:
                logging.exception(traceback.format_exc())
                return JSONResponse(
                    status_code=fastapi.status.HTTP_500_INTERNAL_SERVER_ERROR,
                    content={
                        'msg': traceback.format_exc()
                    }
                )


if __name__ == '__main__':
    uvicorn.run(app=app, host="0.0.0.0", port=30012)
